import random

from common.random_util import select_one_from_weight_map


class RolePool:
    ABSTRACT_5 = 0
    ROLE_5_NORMAL = 1
    ROLE_5_UP = 2
    ABSTRACT_4 = 3
    ABSTRACT_3 = 4
    ROLE_4_NORMAL = 5
    ROLE_4_UP_A = 6
    ROLE_4_UP_B = 7
    ROLE_4_UP_C = 8

    WEIGHT_LIMIT = 10000

    def __init__(self):
        self.not_5_count = 0
        self.not_4_count = 0
        self.must_up = False
        self.must_up_4 = False

    def next(self):
        res = select_one_from_weight_map(self._build_weight_map())
        if res == RolePool.ABSTRACT_5:
            self.not_5_count = 0
            if self.must_up:
                self.must_up = False
                return RolePool.ROLE_5_UP
            if random.random() < 0.5:
                self.must_up = False
                return RolePool.ROLE_5_UP
            else:
                self.must_up = True
                if random.random() < 0.1:
                    return RolePool.ROLE_5_UP
                else:
                    return RolePool.ROLE_5_NORMAL
        self.not_5_count += 1
        if res == RolePool.ABSTRACT_4:
            self.not_4_count = 0
            if self.must_up_4:
                self.must_up_4 = False
                res = select_one_from_weight_map({
                    RolePool.ROLE_4_UP_A: 10,
                    RolePool.ROLE_4_UP_B: 10,
                    RolePool.ROLE_4_UP_C: 10,
                })
            else:
                res = select_one_from_weight_map({
                    RolePool.ROLE_4_UP_A: 10,
                    RolePool.ROLE_4_UP_B: 10,
                    RolePool.ROLE_4_UP_C: 10,
                    RolePool.ROLE_4_NORMAL: 30,
                })
                if res == RolePool.ROLE_4_NORMAL:
                    self.must_up_4 = True
        else:
            self.not_4_count += 1
        return res

    def _build_weight_map(self):
        weight_map = {}
        i = self.not_5_count + 1
        j = self.not_4_count + 1

        weight_5 = 60
        if i >= 74:
            weight_5 += 600 * (i - 73)
        weight_map[RolePool.ABSTRACT_5] = weight_5
        if weight_5 >= RolePool.WEIGHT_LIMIT:
            return weight_map

        weight_4 = 510
        if j >= 9:
            weight_4 += 5100 * (j - 8)
        weight_4 = min(weight_4, RolePool.WEIGHT_LIMIT - weight_5)
        weight_map[RolePool.ABSTRACT_4] = weight_4
        if (weight_4 + weight_5) >= RolePool.WEIGHT_LIMIT:
            return weight_map

        weight_3 = min(9430, RolePool.WEIGHT_LIMIT - weight_5 - weight_4)
        weight_map[RolePool.ABSTRACT_3] = weight_3

        return weight_map
